package com.codestudy.service.impl;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.InputRequiredException;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.codestudy.config.WebSocket.UserContextManager;
import com.codestudy.entity.AIQuestion;
import com.codestudy.service.AIQuestionRepository;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.stereotype.Service;

import java.time.LocalDateTime;

@Service
@Slf4j
public class AIQuestionService {

    @Autowired
    private AIQuestionRepository aiQuestionRepository;

    public String askAI(String question, int userId) throws NoApiKeyException, InputRequiredException {
        // 使用上下文信息生成回答
        Flowable<GenerationResult> result = askAIWithContext(question, userId);
        StringBuilder answerBuilder = new StringBuilder();
        result.blockingForEach(genResult -> {
            String content = genResult.getOutput().getChoices().get(0).getMessage().getContent();
            answerBuilder.append(content);
        });
        return answerBuilder.toString();
    }

    public Flowable<GenerationResult> askAIWithContext(String question, int userId) throws NoApiKeyException, InputRequiredException {
        UserContextManager.ConversationContext context = UserContextManager.getContext(userId);
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(question).build();
        String conversationHistory = context.getConversationHistory();

        GenerationParam param = GenerationParam.builder()
                .apiKey("sk-42f6281d8f4d4ef4be2289fab7707fe6")
                .model("qwen-plus")
                .messages(createMessageListFromHistory(conversationHistory, userMsg))
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .incrementalOutput(true)
                .build();

        return new Generation().streamCall(param).doOnNext(genResult -> {
            String content = genResult.getOutput().getChoices().get(0).getMessage().getContent();
            // 更新上下文
            UserContextManager.addOrUpdateContext(userId, question, content);

            // 保存对话记录到数据库
            saveConversationRecord(question, content, userId);
        });
    }

    private void saveConversationRecord(String question, String answer, int userId) {
        AIQuestion record = new AIQuestion();
        record.setUserId(userId);
        record.setQuestion(question);
        record.setAnswer(answer);
        record.setCreateTime(LocalDateTime.now());
        record.setUpdateTime(LocalDateTime.now());

        aiQuestionRepository.save(record);
    }

    private java.util.List<Message> createMessageListFromHistory(String history, Message newUserMsg) {
        java.util.List<Message> messages = new java.util.ArrayList<>();
        if (!history.isEmpty()) {
            for (String line : history.split("\n")) {
                String[] parts = line.split(": ", 2);
                if (parts.length == 2) {
                    Role role = "User".equals(parts[0]) ? Role.USER : Role.ASSISTANT;
                    messages.add(Message.builder().role(role.getValue()).content(parts[1]).build());
                }
            }
        }
        messages.add(newUserMsg); // 添加当前用户的提问
        return messages;
    }
}